#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <mpi.h>

// find the grid cells covered by particle at x & compute the CIC weights

float dist(float x1[], float x2[]) {
        return sqrt(pow(x1[0]-x2[0], 2.0)+pow(x1[1]-x2[1],2.0)+pow(x1[2]-x2[2],2.0));
}

void TH_weights(float x[], int ng, float *xa, int ind[3][3], float w[3][3])
{
  int i,j,k;
  float dx = xa[1] - xa[0];
  float x2[3];
  float R;
  R=8.0/2400.;

  for (i=0; i<3; i++) {
    // first make sure that positions fall in box volume
    while (x[i] < xa[0]-dx/2) x[i] += ng*dx;
    while (x[i] > xa[ng-1]+dx/2) x[i] -= ng*dx;

    // now figure out nearest grid cells
    ind[i][1] = (int)floor((x[i]-xa[0])/dx);
    if (ind[i][1]<0 || ind[i][1]>=ng-1) {
	ind[i][1]=ng-1; ind[i][2]=0;
}
}
//  x2[0]=xa[ind[0][1]]; x2[1]=xa[ind[1][1]]; x2[2]=xa[ind[2][1]];
//  if (dist(x, x2)<R) {
//	w[0][1]=w[1][1]=w[2][1]=1;
 // } else {
//	w[0][1]=w[1][1]=w[2][1]=0;
 // }
  
    // now search the nearby grid centers
   
    for (i=-1; i<2; i++) {
	for (j=-1; j<2; j++) {
	    for (k=-1; j<2; k++) {
                
	        x2[0]=xa[ind[0][1]+i];x2[1]=xa[ind[1][1]+j];x2[2]=xa[ind[2][1]+k];
		if (dist(x,x2)<R) {w[0][i+1]=1;w[1][j+1]=1;w[2][k+1]=1;} else {
			w[0][i+1]=0;w[1][j+1]=0;w[2][k+1]=0;
		}
	    }
	}
    }
   
  return;
}

FILE *fopen_check(char *filename, char *mode) 
{
  FILE *fp=fopen(filename,mode);
  if (fp==NULL) {
    fprintf(stderr,"Error opening file %s with mode %s\n",filename,mode);
    perror("fopen");
    MPI_Abort(MPI_COMM_WORLD, 1234); exit(EXIT_FAILURE);
  }
  return fp;
}

float ***tensor3(int n1, int n2, int n3)
{
  int i, j;
  float ***m;
  m=(float ***)malloc(n1*sizeof(float **));
  *m=(float **)malloc(n1*n2*sizeof(float *));
  if (!(**m=(float *)calloc(n1*n2*n3,sizeof(float)))) {
    fprintf(stderr, "Error: unable to allocate %d*%d*%d tensor\n",n1,n2,n3);
    MPI_Abort(MPI_COMM_WORLD, 1234); exit(EXIT_FAILURE);
  }
  for(i=0;i<n1;i++) {
    if (i) {m[i]=m[i-1]+n2; *(m[i]) = *(m[i-1]) + n2*n3;}
    for (j=1; j<n2; j++) m[i][j] = m[i][j-1]+n3;
  }
  return m;
}

int main(int argc, char *argv[])
{
  float ***g, *gsum=NULL, *xa, x[3],x2[3], d;
  long n3=0, npart, ng3;
  float w[3][3];
  int f, p, i, j, k, ind[3][3], nfile=128, ng=128;
  int rank, nproc;
  FILE *fp;
  char filename[200];
  double boxsize;
  float R, dx;
  
  MPI_Init(&argc, &argv);
  MPI_Comm_size(MPI_COMM_WORLD, &nproc);
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  
  if (argc < 2 || argc > 4) {
    if (rank==0) 
      fprintf(stderr,"Usage: %s particlefile [outfile] np1\n",argv[0]);
    MPI_Finalize();
    exit(0);
  }

  nfile = nproc; // follow this convention for the IC/gadget/bindelta codes all the way!
  f = rank;

  if (nfile==1) sprintf(filename, "%s", argv[1]);
  else sprintf(filename, "%s.%d", argv[1],f); // previous version had .%02d but for gadget this is the filname naming convention.
  fp=fopen_check(filename, "rb");
 // get p-- the # particles in the current filename  
  fseek(fp, 8, SEEK_SET);
  fread(&p, 4, 1, fp); // this gives the total # particles in this file, checked!
//  printf("particles: %d in file %d \n", p,f);
  npart = p;
  ng = atoi(argv[3]);
  
  fseek(fp, 132, SEEK_SET);
  fread(&boxsize, 8, 1, fp); //gives the correct boxsize---checked.
  boxsize = (float) boxsize; 
  
  MPI_Allreduce(&npart, &n3, 1, MPI_LONG, MPI_SUM, MPI_COMM_WORLD);
  
  xa = calloc(ng,4);
  for (i=0; i<ng; i++) xa[i] = (i+0.5)/ng; dx=xa[1]-xa[0];
  // create & zero the density grid
  g=tensor3(ng,ng,ng);
  ng3 = ng*ng*ng;
  d = ng3*1.0/n3;
  //  printf("boxsize: %g\n", boxsize); 2400 checked!!

  fseek(fp, 268, SEEK_SET);
  while (p>0) {
    --p;
    if (fread(x,4,3,fp) != 3) {
      fprintf(stderr, "Error reading particle from file %s\n", filename);
      perror("fread"); exit(EXIT_FAILURE);
    }
    x[0]/=boxsize;x[1]/=boxsize;x[2]/=boxsize;
 // first make sure that the point lies within the box
    TH_weights(x, ng, xa, ind, w); 
    for (i=-1; i<2; i++) for (j=-2; j<2; j++) for (k=-2; k<2; k++) {
    g[ind[0][1]+i][ind[1][1]+j][ind[2][1]+k]+=d*w[0][i+1]*w[1][j+1]*w[2][k+1];
} 
//    CIC_weights(x,ng,xa,ind,w);
    // now add densities to cells with appropriate weights
  }
  fclose(fp);
  free(xa);

  if (rank==0) gsum = calloc(ng3, sizeof(float));
  MPI_Reduce(**g, gsum, ng3, MPI_FLOAT, MPI_SUM, 0, MPI_COMM_WORLD);
  free(**g); free(*g); free(g);

  MPI_Finalize();
  float totalsum=0.;
  if (rank == 0) {
    for (i=0; i<ng3; i++) totalsum+=gsum[i];
    for (i=0; i<ng3; i++) gsum[i]=gsum[i]/totalsum-1.;
    if (argc==4) fp = fopen_check(argv[argc-2],"wb");
    else fp = stdout;
    fwrite(gsum, sizeof(float), ng3, fp);
    if (argc==4) fclose(fp);
  }

  return 0;
}
